from typing import Set
import os
from .processor import VideoProcessor
import cv2
from skimage.metrics import structural_similarity as ssim


class KeyframeSaver:
    def __init__(self, output_dir: str, ssim_threshold: float):
        self.output_dir = output_dir
        self.ssim_threshold = ssim_threshold
        os.makedirs(output_dir, exist_ok=True)

    def _compare_ssim(self, current, last) -> bool:
        current_gray = cv2.cvtColor(current, cv2.COLOR_BGR2GRAY)
        last_gray = cv2.cvtColor(last, cv2.COLOR_BGR2GRAY)
        score, _ = ssim(current_gray, last_gray, full=True)
        return score < self.ssim_threshold

    def _format_time(self, seconds: float) -> str:
        """将秒数格式化为 HH:MM:SS.mmm 格式"""
        minutes = int((seconds % 3600) // 60)
        seconds = seconds % 60
        return f"{minutes:02d}:{seconds:06.3f}"

    def save(self, video_path: str, frame_ids: Set[int]):
        with VideoProcessor(video_path, 0, 0) as processor:
            fps = processor.fps
            for frame_id, frame in enumerate(processor.get_frames()):
                if frame_id in frame_ids:
                    # 计算时间点
                    time_in_seconds = frame_id / fps
                    time_str = self._format_time(time_in_seconds)
                    # 使用时间命名图片
                    cv2.imwrite(os.path.join(self.output_dir, f"keyframe_{time_str}.jpg"), frame)